from app.agent.llm import get_llm
from app.rag.retriever import Retriever
from app.agent import AgentFactory
from app.graph_rag.graph_builder import graph_builder
from app.graph_rag.state import GraphState
from langchain_core.messages import HumanMessage, AIMessage
import json
# 从core.logger导入get_logger函数
from ..core.logger import get_logger

# 使用get_logger函数初始化logger
logger = get_logger('services.rag_chat_service')

# 定义RAG提示模板
RAG_PROMPT_TEMPLATE = """
你是一个智能知识助手，请根据以下检索内容回答问题。

任务要求：
1. 仔细阅读并分析检索内容
2. 只提取与问题直接相关的信息进行回答
3. 回答要简短明了，直接针对问题给出具体答案
4. 对于数值型问题（如持股数量、百分比等），请直接提供具体数字，避免冗长描述
5. 如果检索内容中没有相关信息或信息不充分，请说"我不知道"
6. 不要添加检索内容中未提及的信息或进行额外解释

检索内容：
{context}

问题：
{question}

回答：
"""

class RAGChatService:
    def __init__(self):
        self.retriever = Retriever()
        self.llm = get_llm()
        # 编译GraphRAG图
        self.graph_rag = graph_builder.build_graph().compile()
    
    def query_company_background(self, query: str) -> str:
        """
        查询公司背景信息
        
        Args:
            query: 用户查询
            
        Returns:
            公司背景信息回答
        """
        logger.info(f"开始查询公司背景信息，查询内容: {query}")
        
        try:
            # 提取公司名称
            company_name = self._extract_company_name(query)
            if not company_name:
                logger.warning("无法从查询中提取公司名称")
                return "抱歉，无法识别查询中涉及的公司名称。"
            
            logger.info(f"提取到公司名称: {company_name}")
            
            # 使用专门的公司背景搜索方法
            search_results = self.retriever.search(
                query=query,
                k=10,
                search_type='basic'
            )
            
            if not search_results:
                logger.warning("未找到相关文档")
                return "抱歉，未找到相关公司背景信息。"
            
            # 构建上下文
            context = "\n\n".join(search_results)
            
            # 构建提示词
            prompt = RAG_PROMPT_TEMPLATE.format(
                context=context,
                question=query
            )
            
            # 调用LLM生成回答
            logger.info("调用LLM生成回答")
            response = self.llm.invoke(prompt)
            answer = response.content if hasattr(response, 'content') else str(response)
            
            logger.info("公司背景信息查询完成")
            return answer
            
        except Exception as e:
            logger.error(f"查询公司背景信息时出错: {e}", exc_info=True)
            return "处理公司背景查询时发生错误。"
    
    def _extract_company_name(self, query: str) -> str:
        """
        从查询中提取公司名称
        
        Args:
            query: 用户查询
            
        Returns:
            公司名称
        """
        import re
        
        # 常见的公司名称模式
        patterns = [
            r'(.+?)股份有限公司',
            r'(.+?)有限公司',
            r'(.+?)集团公司',
            r'(.+?)公司'
        ]
        
        for pattern in patterns:
            match = re.search(pattern, query)
            if match:
                company_name = match.group(1) + match.group(0)[len(match.group(1)):]
                # 进一步清理公司名称
                company_name = re.sub(r'^[^\u4e00-\u9fa5]+', '', company_name)  # 移除开头的非中文字符
                company_name = re.sub(r'[^\u4e00-\u9fa5]+$', '', company_name)  # 移除结尾的非中文字符
                if len(company_name) > 2:  # 确保名称长度合理
                    return company_name
        
        # 如果没有匹配到，尝试提取包含"公司"的最长片段
        company_matches = re.findall(r'[^,，。.!?！？\s]*公司[^,，。.!?！？\s]*', query)
        if company_matches:
            # 返回最长的匹配项
            return max(company_matches, key=len)
        
        return ""
    
    async def generate_stream(self, messages, user_id, conversation_id, 
                              on_complete=None):
        """生成聊天回复流"""
        try:
            # 构建初始状态
            state: GraphState = {
                "messages": messages,
                "user_id": user_id,
                "conversation_id": conversation_id
            }
            
            # 执行GraphRAG图
            final_state = self.graph_rag.invoke(state)
            
            # 获取最终响应 - 修复：从response字段而不是messages字段获取
            response_content = final_state.get("response", "")
            if not response_content:
                # 如果response字段为空，回退到原来的方式
                response_messages = final_state.get("messages", [])
                if response_messages:
                    response_content = response_messages[-1].content if hasattr(response_messages[-1], 'content') else str(response_messages[-1])
                else:
                    response_content = "抱歉，无法生成回复。"
            
            # 保存用户消息到数据库
            if on_complete and messages:
                user_message = messages[-1]["content"] if isinstance(messages[-1], dict) else str(messages[-1])
                await on_complete(conversation_id=conversation_id, 
                                  content=user_message, 
                                  role="user")
            
            # 模拟流式输出，逐字符返回
            for char in response_content:
                yield char
                
            # 保存AI回复到数据库
            if on_complete:
                await on_complete(conversation_id=conversation_id, 
                                  content=response_content, 
                                  role="assistant")
            
            logger.info(f"Chat completed with response: {response_content}")
            
        except Exception as e:
            logger.error(f"Error in generate_stream: {str(e)}")
            error_message = "抱歉，处理您的请求时出现错误。"
            if on_complete:
                # 保存用户消息
                if messages:
                    user_message = messages[-1]["content"] if isinstance(messages[-1], dict) else str(messages[-1])
                    await on_complete(conversation_id=conversation_id, 
                                      content=user_message, 
                                      role="user")
                
                # 保存错误消息
                await on_complete(conversation_id=conversation_id, 
                                  content=error_message, 
                                  role="assistant")
            for char in error_message:
                yield char

    async def chat(self, messages: list, user_id: int, conversation_id: int, 
                   on_complete=None):
        """处理聊天请求"""
        logger.info(f"Starting chat with {len(messages)} messages")
        
        try:
            # 构建初始状态
            state: GraphState = {
                "messages": messages,
                "user_id": user_id,
                "conversation_id": conversation_id
            }
            
            # 执行GraphRAG图
            final_state = self.graph_rag.invoke(state)
            
            # 获取最终响应 - 修复：从response字段而不是重新执行获取
            full_response = final_state.get("response", "")
            if not full_response:
                # 如果response字段为空，回退到原来的方式
                response_messages = final_state.get("messages", [])
                if response_messages:
                    full_response = response_messages[-1].content if hasattr(response_messages[-1], 'content') else str(response_messages[-1])
                else:
                    full_response = "抱歉，无法生成回复。"
            
            # 保存对话历史
            if on_complete:
                # 保存用户消息
                if messages:
                    user_message = messages[-1]["content"] if isinstance(messages[-1], dict) else str(messages[-1])
                    await on_complete(conversation_id=conversation_id, 
                                      content=user_message, 
                                      role="user")
                
                # 保存AI响应
                await on_complete(conversation_id=conversation_id, 
                                  content=full_response, 
                                  role="assistant")
            
            logger.info(f"Chat completed with response: {full_response}")
            return full_response
            
        except Exception as e:
            logger.error(f"Error in chat: {str(e)}")
            error_message = "抱歉，处理您的请求时出现错误。"
            # 保存错误消息到数据库
            if on_complete:
                # 保存用户消息
                if messages:
                    user_message = messages[-1]["content"] if isinstance(messages[-1], dict) else str(messages[-1])
                    await on_complete(conversation_id=conversation_id, 
                                      content=user_message, 
                                      role="user")
                
                # 保存错误消息
                await on_complete(conversation_id=conversation_id, 
                                  content=error_message, 
                                  role="assistant")
            return error_message


class ComplexFinanceQuery:
    def execute_query(self, query_text):
        # 导入真正的ComplexFinanceQuery类
        from app.agent.complex_finance_query import ComplexFinanceQuery as RealComplexFinanceQuery
        
        # 创建真正的ComplexFinanceQuery实例
        real_query_handler = RealComplexFinanceQuery()
        
        # 调用process_query方法（原execute_query的逻辑已整合到process_query中）
        return real_query_handler.process_query(query_text)


class QueryHandler:
    def __init__(self):
        self.rag_service = RAGChatService()
    
    def handle_query(self, query_text, intent_type):
        # 根据意图类型处理查询
        if intent_type == "complex_finance_query":
            logger.info("处理复杂金融查询意图")
            complex_finance_handler = ComplexFinanceQuery()
            try:
                # 尝试处理复杂金融查询
                result = complex_finance_handler.execute_query(query_text)
                # 检查是否是回退消息
                if "QUERY_REQUIRES_RAG" in result or "已转交相关模块处理" in result:
                    logger.info("复杂金融查询无法处理，转交company_background处理")
                    intent_type = "company_background"
                    # 继续下面的company_background处理流程
                else:
                    logger.info("复杂金融查询处理完成")
                    return result
            except Exception as e:
                logger.error(f"处理复杂金融查询时出错: {e}")
                # 出错时也尝试转到company_background处理
                intent_type = "company_background"
        
        # 处理公司背景查询
        if intent_type == "company_background":
            logger.info("处理公司背景查询意图")
            try:
                # 使用GraphRAG处理公司背景查询
                # 创建初始状态
                initial_state = GraphState(
                    messages=[HumanMessage(content=query_text)],
                    processing_type="prospectus_rag",  # 直接设置为prospectus_rag类型
                    intent_data={"intent": "company_background", "entities": {}, "confidence": 0.95}
                )
                
                # 调用GraphRAG流程
                final_state = self.rag_service.graph_rag.invoke(initial_state)
                
                # 从最终状态中提取响应
                response = final_state.get("response", "")
                if response and response.strip():
                    logger.info("公司背景查询处理完成")
                    # 确保返回的是实际的响应内容而不是原始查询
                    if response.strip() != query_text.strip():
                        return response
                    else:
                        # 如果响应和查询相同，尝试从消息历史中获取AI回复
                        messages = final_state.get("messages", [])
                        if messages and len(messages) >= 2:
                            last_message = messages[-1]
                            if hasattr(last_message, 'content') and last_message.content != query_text:
                                return last_message.content
                        return "抱歉，未找到相关公司背景信息。"
                else:
                    logger.warning("公司背景查询未返回有效结果")
                    return "抱歉，未找到相关公司背景信息。"
            except Exception as e:
                logger.error(f"处理公司背景查询时出错: {e}")
                return "处理公司背景查询时发生错误。"
